package main

import (
	"fmt"
	"github.com/pebbe/zmq4"
	"log"
	// "math"
	// "math/rand"
	"time"
	// "time"
)

var _ = fmt.Println

func ce(err error) {
	if err != nil {
		log.Panic(err)
	}
}

// func newSocks(t zmq4.Type, endpoint string) (socket *zmq4.Socket, err error) {
// 	context, err := zmq4.NewContext()
// 	if err != nil {
// 		return nil, err
// 	}
// 	socket, err = context.NewSocket(t)
// 	if err != nil {
// 		return nil, err
// 	}
// 	err = socket.Connect(endpoint)
// 	if err != nil {
// 		socket.Close()
// 		return nil, err
// 	}
// 	return socket, nil
// }

func rrclient(id int) {
	context, err := zmq4.NewContext()
	ce(err)
	requester, err := context.NewSocket(zmq4.REQ)
	ce(err)
	err = requester.Connect("tcp://0.0.0.0:5559")
	ce(err)
	control, err := context.NewSocket(zmq4.PUB)
	ce(err)
	err = control.Bind("tcp://0.0.0.0:5561")
	ce(err)
	for i := 0; i < 3; i++ {
		_, err = requester.SendMessage(fmt.Sprintf("hello from %d", id))
		ce(err)
		res, err := requester.RecvMessage(zmq4.SNDMORE)
		ce(err)
		log.Printf("rrclient %d received %d %s\n", id, i, res)
	}
	log.Println("a")
	_, err = control.SendMessage("EXIT")
	log.Println("b")
	ce(err)
}

func rrworker(id int) {
	context, err := zmq4.NewContext()
	ce(err)
	responder, err := context.NewSocket(zmq4.REP)
	ce(err)
	err = responder.Connect("tcp://0.0.0.0:5560")
	ce(err)

	control, err := context.NewSocket(zmq4.SUB)
	ce(err)
	err = control.Connect("tcp://0.0.0.0:5561")
	ce(err)

	poller := zmq4.NewPoller()
	poller.Add(responder, zmq4.POLLIN)
	poller.Add(control, zmq4.POLLIN)

	for {
		sockets, _ := poller.Poll(-1)
		for _, socket := range sockets {
			switch s := socket.Socket; s {
			case responder:
				res, err := responder.RecvMessage(zmq4.SNDMORE)
				ce(err)
				log.Printf("rrworker %d received %s\n", id, res)
				time.Sleep(time.Second)
				_, err = responder.SendMessage(fmt.Sprintf("%s world %d", res, id))
				ce(err)
			case control:
				log.Println("recv exit")
				goto EXIT
			default:
				panic("")
			}
		}
	}
EXIT:
}

func rrbroker() {
	context, err := zmq4.NewContext()
	ce(err)

	frontend, err := context.NewSocket(zmq4.ROUTER)
	ce(err)
	err = frontend.Bind("tcp://0.0.0.0:5559")
	ce(err)

	backend, err := context.NewSocket(zmq4.DEALER)
	ce(err)
	err = backend.Bind("tcp://0.0.0.0:5560")
	ce(err)

	poller := zmq4.NewPoller()
	poller.Add(frontend, zmq4.POLLIN)
	poller.Add(backend, zmq4.POLLIN)
	for {
		sockets, _ := poller.Poll(-1)
		for _, socket := range sockets {
			switch s := socket.Socket; s {
			case frontend:
				res, err := frontend.RecvMessage(zmq4.SNDMORE)
				ce(err)
				_, err = backend.SendMessage(res)
				ce(err)
			case backend:
				res, err := backend.RecvMessage(zmq4.SNDMORE)
				ce(err)
				_, err = frontend.SendMessage(res)
				ce(err)
			}
		}
	}
}

func msgqueue() {
	context, err := zmq4.NewContext()
	ce(err)

	frontend, err := context.NewSocket(zmq4.ROUTER)
	ce(err)
	err = frontend.Bind("tcp://0.0.0.0:5559")
	ce(err)

	backend, err := context.NewSocket(zmq4.DEALER)
	ce(err)
	err = backend.Bind("tcp://0.0.0.0:5560")
	ce(err)

	zmq4.Proxy(frontend, backend, nil)
}

func main() {
	log.SetFlags(log.Lshortfile)
	for i := 0; i < 1; i++ {
		go rrworker(i)
	}
	for i := 0; i < 1; i++ {
		go rrclient(i)
	}
	// rrbroker()
	msgqueue()
}
